#!/usr/bin/env python
# -*- coding: utf-8 -*-

import re
import os
import os.path
import wave
import contextlib
import shutil
import subprocess

from mutagen.mp3 import MP3
from mutagen.id3 import TIT2, TALB, TPE1, TRCK, TDRC, TCON, COMM
from mutagen.flac import FLAC
from mutagen.mp4 import MP4
from mutagen.asf import ASF
from mutagen.musepack import Musepack
from mutagen.oggvorbis import OggVorbis
from mutagen.oggflac import OggFLAC

from pydub import AudioSegment

from album import Album
from musicbrainz import MusicBrainz
from musicbrainz import MBError
from bagofwords import BagOfWords
from observable import Observable

class SimilarityError(Exception): pass

class Kelindo(Observable):

    CODES = ['OK', 'NO_MEDIA', 'TOOMANY_MEDIA', 'NO_TNUMBERS', 'NO_DURATIONS', 'UNSUPPORTED', 'NO_MATCH', 'MATCH']

    def __init__(self, path):
        super(Kelindo, self).__init__()
        
        self.dir = path
        self.__invalidExts = set()
        self.__validExts = ['.mp3', '.flac', '.wav', '.ogg', '.mp4', '.m4a', '.asf', '.mpc']
        
        self.albums = []
        
        self.mb = MusicBrainz()
        self.stat_start = None
        self.stat_stop = None
        
                
    def __coherent_track_position(self, album):
        """ Checks if track numbers within an album are coherent """
        album.tracks = sorted(album.tracks, key=lambda track: track['trackposition'])
        for i in range(len(album.tracks)):
            if album.tracks[i]['trackposition'] != i + 1:
                return False
        return True
        
        
    def __build_album(self, root, files):
        """ Roughly classifies dir-tree """     
        album = Album(root) 
        exts = set()

        for file in files:
            trackName, ext = os.path.splitext(file)
            if ext.lower() in self.__validExts:
                exts.add(ext)
                track = {
                    'filename' : file,
                    'trackposition' : 0,
                }
                album.tracks.append(track)
                
                # search track numbers
                match = re.search('(\d{1,2})', file) 
                if match:
                    track['trackposition'] = int(match.groups()[0])
            # debug only: annotate unrecognized extensions
            else:
                self.__invalidExts.add(ext)
                
        if len(exts) == 1:
            album.format = exts.pop().lower()
            album.format = album.format.replace('.','')
        else:
            album.format = 'misc'
                
        if (len(album.tracks) < 1):
            album.status = 'NO_MEDIA'
            return album
        
        if (len(album.tracks) > 99):
            album.status = 'TOOMANY_MEDIA'
            return album
        
        try:
            if not self.__coherent_track_position(album):
                self.__get_metadata(album)
            else:
                self.__get_metadata(album, ['TALB', 'TIT2', 'TPE1', 'TCON'])
            
            if not self.__coherent_track_position(album):
                album.status = 'NO_TNUMBERS'
                return album
        except:
            album.status = 'UNSUPPORTED'
            return album
                
        album.status = 'OK'
        return album
    
    
    def __get_metadata(self, album, tags = ['TALB', 'TIT2', 'TPE1', 'TCON', 'TRCK']):
        """ Extracts metadata for all tracks of an album """
        tagMap = {
                'TALB': 'albumtitle',
                'TIT2': 'tracktitle',
                'TPE1': 'trackauthors',
                'TCON': 'trackgenre',
                'TRCK': 'trackposition',
        }

        for track in album.tracks:
            filename = os.path.join(album.dir, track['filename'])
            extension = os.path.splitext(track['filename'])[1].lower()
            
            if extension == '.mp3':
                meta = MP3(filename)
            elif extension == '.flac':
                meta = FLAC(filename)
            elif extension == '.mp4' or extension.lower() == '.m4a':
                meta = MP4(filename)
            elif extension == '.asf':
                meta = ASF(filename)
            elif extension == '.mpc':
                meta = Musepack(filename)
            elif extension == '.ogg':
                meta = Ogg(filename)
            elif extension == '.wav':
                album.need_conversion = True
                with contextlib.closing(wave.open(filename,'r')) as wav:
                    track['tracklength'] = wav.getnframes() / wav.getframerate()
            else:
                raise IOError
                
            try:
                track['tracklength'] = int(meta.info.length) 
                track['trackbitrate'] = (float(os.path.getsize(filename) * 8.0) / 1000.0) / track['tracklength']
            except AttributeError:
                pass

            for tag in tags:
                try:
                    track[tagMap[tag]] = meta[tag].text[0]
                except KeyError:
                    pass
                    
            if 'TRCK' in tags:
                track['trackposition'] = int(unicode(track['trackposition']).split('/')[0])

        album.avg_bitrate = 0.0
        for track in album.tracks:
            album.avg_bitrate += track['trackbitrate']
        album.avg_bitrate /= len(album.tracks)
        return
                
                
    def __get_best_match(self, album, candidates):
        """ Returns the best match among a set of candidates """
        album_bag = BagOfWords()
        for word in album.dir.split('/')[-3:]: album_bag.add_string(word)
        for track in album.tracks:
            for field in ['albumtitle', 'tracktitle', 'trackauthors', 'filename']:
                try:
                    album_bag.add_string(track[field])
                except KeyError:
                    pass
            
        bestMatch = [0.0, None]
        for candidate in candidates:
            if len(album.tracks) != len(candidate):
                continue
            
            other_bag = BagOfWords()
            for track in candidate:
                for field in ['albumtitle', 'albumauthors', 'tracktitle', 'trackauthors']:
                    other_bag.add_string(track[field])

            similarity = album_bag.similarity(other_bag)
            if (similarity > bestMatch[0]):
                bestMatch[0] = similarity
                bestMatch[1] = candidate
                
        if bestMatch[0] == 0.0:
            raise SimilarityError
        return bestMatch      
            
    def __classify_album(self, album): 
        #TODO: get_MBID and get_albums might fail
        # we still do not handle that case
        MBIDs = self.mb.get_MBID(album.get_durations())
        if len(MBIDs) > 15:
            album.status = 'NO_MATCH'
            return
        try:    
            candidates = []
            for MBID in MBIDs:
                candidate = self.mb.get_albums(MBID)
                candidates.extend(candidate)
            album.distance, album.best_candidate = self.__get_best_match(album, candidates)
            album.status = 'MATCH'
            album.finalize = True
        
        except MBError:
            album.status = 'NO_MATCH'
            return
        except SimilarityError:
            album.status = 'NO_MATCH'
            return
        return
        
    def __write_tags(self, dst, track, ext):
        if ext == '.mp3':
            audio = MP3(dst)
            audio['TIT2'] = TIT2(encoding=3, text=[track['tracktitle']])
            audio['TALB'] = TALB(encoding=3, text=[track['albumtitle']])
            audio['TRCK'] = TRCK(encoding=3, text=[track['trackposition']])
            audio['TPE1'] = TPE1(encoding=3, text=[track['albumauthors']])
            audio['TDRC'] = TDRC(encoding=3, text=[track['albumyear']])
            audio['COMM'] = COMM(encoding=3, text=['Beautyfied by Kelindo!'])
        else:
            if ext == '.flac':
                audio = FLAC(dst)
            elif ext == '.mp4' or ext.lower() == '.m4a':
                audio = MP4(dst)
            elif ext == '.asf':
                audio = ASF(dst)
            elif ext == '.mpc':
                audio = Musepack(dst)
            elif ext == '.ogg':
                try:
                    audio = OggVorbis(dst)
                except:
                    print 'no vorbis'
                    audio = OggFLAC(dst)
                print audio
            else:
                pass
            
            audio['title'] = track['tracktitle']
            audio['album'] = track['albumtitle']
            audio['tracknumber'] = track['trackposition']
            audio['artist'] = track['albumauthors']       
            audio['date'] = track['albumyear']
            audio['comment'] = 'Beautyfied by Kelindo!'
        audio.save()
        return
        
        
    def __convert(self, src, dst, fmt, avg_bitrate):
        def closest(target, collection) :
            return min((abs(target - i), i) for i in collection)[1]
        #TODO - Gestire la conversione con piu parametri (ffmpeg)
        bitrate = closest(avg_bitrate, [128, 256, 320])        
        if fmt == 'mp3':
            #subprocess.call(['/usr/bin/sox', '--multi-threaded', src, '--compression', '%d.2' % (bitrate,), dst])
            AudioSegment.from_file(src).export(dst, format=fmt)
        elif fmt == 'ogg':
            #ogg_compression = {128: '5', 256 : '7', 320 : '9'}
            #subprocess.call(['sox', '--multi-threaded', src, '--compression', ogg_compression[bitrate], dst])
            AudioSegment.from_file(src).export(dst, format=fmt)
            
        else:
            #subprocess.call(['sox', '--multi-threaded', src, dst])
            AudioSegment.from_file(src).export(dst, format=fmt)
            
        return
    
    def __copy(self, src, dst):
        shutil.copyfile(src, dst)
        return
    
    def __move(self, src, dst):
        shutil.move(src, dst)
        return
        
    def find_candidate_albums(self):
        self.fire(key='candidate_started')
        for root, dirs, files in os.walk(self.dir):
            album = self.__build_album(root, files) 
            self.albums.append(album)
            self.fire(key='candidate_found')
            
        self.stat_start = self.get_stats()

        if self.stat_start['OK'] == 0:
            self.fire(key='candidate_no')
        else:
            self.fire(key='candidate_done')
        return
        
        
    def classify_albums(self):
        print 'classify_albums'
        album_total = self.get_stats()['OK']
        album_done = 0
        self.fire(key='album_started', album=None, total=album_total, done=album_done)
        for album in self.albums:
            if album.status == 'OK':
                print 'start classify album'
                self.__classify_album(album)
                if album.status == 'MATCH':
                    album_done += 1
                    self.fire(key='album_found', album=album, total=album_total, done=album_done)
        
        self.stat_stop = self.get_stats()
        self.fire(key='album_done', album=None, total=album_total, done=album_done)
        return

    def find_candidate_only_conversion(self):
        print 'TODO'
        return
        
    def finalize_albums(self, root_dst_dir, fmt, move_folder):
        root_dst_dir = os.path.join(root_dst_dir, "Kelindo")

        tracks_total = tracks_done = 0
        
        for album in self.albums:
            if album.finalize == True:                
                if fmt == 'as source':
                    if album.format == 'misc':
                        album.new_format = 'mp3'
                        album.need_conversion = True
                    elif album.format == 'wav':
                        album.new_format = 'flac'
                        album.need_conversion = True
                else:
                    if not album.format == fmt:
                        album.new_format = fmt
                        album.need_conversion = True    
                tracks_total += album.get_track_number()
        
        self.fire(key='finalize_started', album=None, total=tracks_total, done=tracks_done)
        
        for album in self.albums:
            if album.status != 'MATCH' or album.finalize != True:
                continue
            
            dst_dir = os.path.join(root_dst_dir, album.best_candidate[0]['albumauthors'])
            if int(album.best_candidate[0]['trackmedium']) != 1:
                medium = ' (CD %s)' % album.best_candidate[0]['trackmedium']
                dst_dir = os.path.join(dst_dir, album.best_candidate[0]['albumtitle'] + medium)
            else:
                dst_dir = os.path.join(dst_dir, album.best_candidate[0]['albumtitle'])
            
            if not os.path.exists(dst_dir):
                os.makedirs(dst_dir)
                album.new_dir = dst_dir
            else:
                print 'Album already exists, overwrite'
                shutil.rmtree(dst_dir)
                os.makedirs(dst_dir)
                album.new_dir = dst_dir

            for idx in range(1, album.get_track_number() + 1):                
                src_filename, src_ext = os.path.splitext(album.get_track(idx)['filename'])
                src = os.path.join(album.dir, src_filename + src_ext)  
                                
                dst_track = album.get_best_candidate_track(idx)
                dst_filename = '%s - %s' % (
                    dst_track['trackposition'].zfill(2),
                    dst_track['tracktitle'])
                dst = os.path.join(dst_dir, dst_filename)
                
                if album.need_conversion:
                    dst = dst + '.' + fmt
                    self.__convert(src, dst, fmt, album.avg_bitrate)
                    self.__write_tags(dst, dst_track, '.' + fmt)
                else:
                    dst = dst + src_ext
                    if move_folder: 
                        self.__move(src, dst)
                        os.chmod(dst, 0644)
                    else:
                        self.__copy(src, dst)
                        os.chmod(dst, 0644)
                    self.__write_tags(dst, dst_track, src_ext.lower())
                
                tracks_done += 1
                self.fire(key='finalize_track_done', album=album, total=tracks_total, done=tracks_done)
                
            if move_folder:
                shutil.rmtree(album.dir)
          
        self.fire(key='finalize_done', album=None, total=tracks_total, done=tracks_done)
        return
        
        
    def get_stats(self):
        """ Returns some stats """
        stat = {}
        for status in Kelindo.CODES:
            stat[status] = 0
        
        for album in self.albums:
            stat[album.status] += 1
        return stat
        

        
        
